# -*- coding: utf-8 -*-

from .img import *


def _get_size(data):
    i = 2
    ld = len(data) - 2
    while i < ld:
        j = i + 3
        if j < ld:
            lc = int.from_bytes(data[j - 1:j + 1], byteorder='big') + 2
            ni = i + lc
            if ni <= ld:
                chunk = data[i:ni]
                if bytes(chunk[:2]) == b'\xff\xc0':
                    h = int.from_bytes(chunk[5:7], byteorder='big')
                    w = int.from_bytes(chunk[7:9], byteorder='big')
                    return w, h
            else:
                return None, None
            i = ni
        else:
            return None, None
    return None, None


class Jpg(Img):
    __FORMAT__ = 'jpg'
    __MAGIC__ = bytearray(b'\xff\xd8\xff')

    @classmethod
    def size(cls, fobj):
        rsize = 512
        force_stop = False
        while True:
            if rsize > fobj.file_size:
                rsize = fobj.file_size
                force_stop = True
            rb = fobj.read(rsize)
            # 在 rb 中查找 ffc0 数据段
            w, h = _get_size(rb)
            if w:
                return ImgSize(fobj.file_size, w, h)
            if force_stop:
                return None
            rsize += 512

        return None
